#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pycurl
import StringIO
import sys
import json
import os

DATE_START = "1700"
DATE_END = "2100"
SPLIT = 10000


def GetNumberString():
    url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
    buf = StringIO.StringIO()
    c = pycurl.Curl()
    c.setopt(pycurl.POST, 1)
    c.setopt(c.URL, url)
    c.setopt(pycurl.HTTPPOST, [('db', 'pubmed'), ('mindate', DATE_START), ('maxdate', DATE_END), ('retmode', "json")])
    # c.setopt(c.WRITEDATA, buf)
    c.setopt(c.WRITEFUNCTION, buf.write)
    c.perform()
    c.close()
    body = buf.getvalue()
    buf.close()
    return body


def GetNumber(numberString):
    import json
    s = json.loads(numberString)
    return s["esearchresult"]["count"]


def GetPubMed(i):
    try:
        url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
        retstart = str(i * SPLIT)
        buf = StringIO.StringIO()
        c = pycurl.Curl()
        c.setopt(pycurl.POST, 1)
        c.setopt(c.URL, url)
        c.setopt(pycurl.HTTPPOST,
                 [('db', 'pubmed'), ('mindate', DATE_START), ('maxdate', DATE_END), ('retmode', "json"),
                  ('retstart', retstart), ('retmax', str(SPLIT))])
        c.setopt(c.WRITEFUNCTION, buf.write)
        c.perform()
        result = buf.getvalue()
        json.loads(result)
        outfile = open("./list/pmcmeta_list_" + str(i) + ".json", "w")
        outfile.write(buf.getvalue())
        outfile.close()
        buf.close()
        print "PubMed meta task id " + str(i) + " has been successfully downloaded."
    except ValueError:
        os.system("rm -f ./list/pmcmeta_list_" + str(i) + ".json")
        GetPubMed(i)
    except pycurl.error:
        os.system("rm -f ./list/pmcmeta_list_" + str(i) + ".json")
        GetPubMed(i)


if __name__ == '__main__':
    if len(sys.argv) >= 3:
        DATE_START = sys.argv[1]
        DATE_END = sys.argv[2]
    if len(sys.argv) >= 4:
        SPLIT = int(sys.argv[3])
    numberString = GetNumberString()
    number = int(GetNumber(numberString))
    if number != None:
        end = int(number / SPLIT)
        if number % SPLIT != 0:
            end += 1
        for i in xrange(517, end):
            GetPubMed(i)
